# create a task to generate a house
# procthor house to the ai2thor 

import sys,os,json
import copy
import random
from tqdm import tqdm
import numpy as np

from pathlib import Path
from PIL import Image
from matplotlib import pyplot as plt    
import networkx as nx

sys.path.append(str(Path(__file__).parent.parent))

import prior
from utils.util import read_house_from_json,read_json,read_yaml_to_array
from utils.debug import enable_remote_debug
from utils.expert import ShortestPathNavigatorTHOR

from ai2thor.controller import Controller
from ai2thor.platform import CloudRendering
from ai2thor.util.metrics import get_shortest_path_to_object


from third_party.procthor.constants import PROCTHOR_INITIALIZATION
from third_party.procthor.generation import HouseGenerator, PROCTHOR10K_ROOM_SPEC_SAMPLER

from utils.visualization_utils import get_agent_map_data, visualize_agent_path


DEBUG = True
OBJECT_TYPES = read_yaml_to_array("config/habitat2022.yaml")

class PathNavigator:
    def __init__(self, scene, grid_size=0.25, include_move_left_right=False):
        self.controller = Controller(
            scene=scene,
            platform=CloudRendering,
            renderDepthImage=True,
            renderInstanceSegmentation=True,
        )
        self.navigator = ShortestPathNavigatorTHOR(
            self.controller,
            grid_size=grid_size,
            include_move_left_right=include_move_left_right
        )

        self.shortest_path = None
        self.target_object = None
        self.agent_initial = None

    def get_shortest_path_to_object(self, target_object):

        self.agent = self.navigator.controller.last_event.metadata["agent"]
        # save the initial agent position   
        self.agent_initial = copy.deepcopy(self.agent)
        self.target_object = copy.deepcopy(target_object)

        reachable_positions = self.navigator.controller.step(
            action="GetReachablePositions",
        ).metadata["actionReturn"]

        # 获取目标物体的可交互位置 返回对应物体的可交互位置
        target_object_interactable_positions = self.navigator.controller.step(
            action="GetInteractablePoses",
            objectId=target_object['objectId'],
            positions=reachable_positions,
        ).metadata["actionReturn"]
        # 获取目标物体的图节点键

        self.target_keys = self.navigator.get_key(target_object_interactable_positions[0])
        source_key = self.navigator.get_key(self.agent)

        # 计算最短路径
        try:
            self.shortest_path = self.navigator.shortest_state_path(source_key, self.target_keys)
            return self.shortest_path
        except nx.NetworkXNoPath:
            print("cant find the path")
            return None
    
    def inverse_get_key(self):
        """
        AgentLocKeyType ->Dict[str, Any]
        """
        self.positions = []
        for path in self.shortest_path:
            position = {
                "x": path[0],   
                "z": path[1],
                "y": self.controller.last_event.metadata["agent"]["position"]["y"], 
                "rotation": path[2],
                "cameraHorizon": path[3]
            }
            self.positions.append(position)

        return self.positions   
    
    def get_non_wall_objects(self):
        non_wall_objects = [
            obj for obj in self.controller.last_event.metadata["objects"]
            if "wall" not in obj["objectType"] or "wall" not in obj["name"]  # 假设墙的类型包含 "Wall"
        ]
        return non_wall_objects

    def get_shortest_path_action_sequence(self):
        # TODO:check the source and target key
        source_state_key = self.navigator.get_key(self.agent)
        action_sequence = []

        for target_sate_position in self.shortest_path:
            action =  self.navigator.action_transitioning_between_keys(source_state_key, target_sate_position)
            source_state_key = target_sate_position
            action_sequence.append(action)
        return action_sequence
    
    def get_shortest_path_action_sequence_frame(self):
        """
        get the action sequence frame from the action sequence
        x = input_dict["position"]["x"]
        z = input_dict["position"]["z"]
        rot = input_dict["rotation"]["y"]
        hor = input_dict["cameraHorizon"]
        (x,y,z,rot,hor)
        """
        frames = []
        y_position = self.controller.last_event.metadata["agent"]["position"]["y"]
        for position in self.shortest_path:
            self.controller.step(
                action="Teleport",
                position=dict(x=position[0], y=y_position, z=position[1]),
                rotation=dict(x=0 , y=position[2], z=0),
                horizon=position[3],
                standing=True
            )
            frame = self.controller.last_event.frame
            frames.append(frame)
        return frames

    def get_top_down_path_view(self,target_ids):
        paths = []
        for path in self.shortest_path:
            position = {
            "x": path[0],   
            "z": path[1],
            "y": self.controller.last_event.metadata["agent"]["position"]["y"], 
            }
            paths.append(position)
        
        top_down = self.navigator.controller.get_top_down_path_view(paths, target_ids)
        
        return Image.fromarray(top_down)


    def look_at_top_down_path_view(self):
        positions = self.inverse_get_key()
        for i in range(len(positions)):
            teleport_action = {
                "action": "TeleportFull",
                "position": dict(x=positions[i]["x"], y=positions[i]["y"], z=positions[i]["z"]),
                "rotation": dict(x=0 , y=positions[i]["rotation"], z=0),
                "horizon": positions[i]["cameraHorizon"],
                "standing": True
            }
            self.controller.step(teleport_action)
            meta = get_agent_map_data(self.controller)
            frame = visualize_agent_path(
                positions[: i + 1],
                meta["frame"],
                meta["pos_translator"],
                color_pair_ind=2,
                opacity=0,
                only_show_last_visibility_cone=True,
                max_colors=len(positions),
            )
            img = Image.fromarray(frame)
            img.save(f"debug/top_down_path_view/top_{i}.png")  

        # save the movie
        # os.system("ffmpeg -framerate 10 -i debug/top_down_path_view/top_%d.png -c:v libx264 -r 30 -pix_fmt yuv420p debug/top_down_path_view.mp4")
        return

        


    def get_top_down_frame(self):
        # Setup the top-down camera
        event = self.controller.step(action="GetMapViewCameraProperties", raise_for_failure=True)
        pose = copy.deepcopy(event.metadata["actionReturn"])

        bounds = event.metadata["sceneBounds"]["size"]
        max_bound = max(bounds["x"], bounds["z"])

        pose["fieldOfView"] = 50
        pose["position"]["y"] += 1.1 * max_bound
        pose["orthographic"] = False
        pose["farClippingPlane"] = 50
        del pose["orthographicSize"]

        # add the camera to the scene
        event = self.controller.step(
            action="AddThirdPartyCamera",
            **pose,
            skyboxColor="white",
            raise_for_failure=True,
        )
        top_down_frame = event.third_party_camera_frames[-1]
        return Image.fromarray(top_down_frame)


    def visualize_reachable_positions(self):
        """
        Visualizes the reachable positions in the scene using the given controller.

        Parameters:
        - controller: The AI2-THOR controller object.
        """
        # Get the event from the controller
        event = self.controller.step(action="GetReachablePositions")
        
        # Extract reachable positions
        reachable_positions = event.metadata["actionReturn"]
        xs = [rp["x"] for rp in reachable_positions]
        zs = [rp["z"] for rp in reachable_positions]

        # Plot the reachable positions
        fig, ax = plt.subplots(1, 1)
        ax.scatter(xs, zs)
        ax.set_xlabel("$x$")
        ax.set_ylabel("$z$")
        ax.set_title("Reachable Positions in the Scene")
        ax.set_aspect("equal")
        plt.savefig("debug/reachable_positions.png")

        return fig
    
    def save_navigation_data(self, file_path="debug/navigation_data.npy"):
        actions = self.get_shortest_path_action_sequence()
        frames = self.get_shortest_path_action_sequence_frame()
        data = {
            'object_id': self.target_object,
            'agent_initial_position': self.agent_initial,
            'shortest_path': self.inverse_get_key(),
            'actions': actions,
            'frames': frames
        }
        np.save(file_path, data)


if __name__=="__main__":
    # enable_remote_debug()
    # debugpy.breakpoint()
    scene = read_json("big-dataset/house.json")
    path_navigator = PathNavigator(scene=scene)

    non_wall_objects = path_navigator.get_non_wall_objects()
    
    target_object = random.choice(non_wall_objects)
    print(target_object['objectId'])
    path = path_navigator.get_shortest_path_to_object(target_object)
    print(path)

    top_down_path_view = path_navigator.look_at_top_down_path_view()
    # top_down_path_view.save("debug/top_down_path_view.png")

    top_down_path_view = path_navigator.get_top_down_path_view(target_object['objectId'])
    top_down_path_view.save("debug/top_down_path_view.png")

    top_down_frame = path_navigator.get_top_down_frame()
    top_down_frame.save("debug/top_down_frame.png")

    action_sequence = path_navigator.get_shortest_path_action_sequence()
    print(action_sequence)

    path_navigator.visualize_reachable_positions()

    frames = path_navigator.get_shortest_path_action_sequence_frame()
    for i, frame in enumerate(frames):
        image = Image.fromarray(frame)
        image.save(f"debug/frame_sequence/frame_{i}.png")  

    os.system("ffmpeg -framerate 10 -i debug/frame_sequence/frame_%d.png -c:v libx264 -r 30 -pix_fmt yuv420p debug/frame_sequence.mp4") 
    path_navigator.save_navigation_data()


